import pandas as pd
import numpy as np
from sklearn.ensemble import RandomForestRegressor
import joblib
from datetime import datetime, timedelta


class SalesPredictor:
    def __init__(self, model_path='sales_model.pkl'):
        self.model_path = model_path
        self.model = None
        self.cutoff_date = None
        self.max_days = 90  # 商品生命周期90天
        self.forecast_days = 14  # 预测未来14天

    def prepare_data(self, raw_df, cutoff_date=None):
        """
        准备训练数据：
        1. 过滤超过90天的数据
        2. 补充缺失日期销量为0
        3. 生成特征和目标变量
        """
        # 确保日期格式为YYYY-MM-DD
        raw_df['date_id'] = pd.to_datetime(raw_df['date_id'], format='%Y-%m-%d')
        raw_df['new_date'] = pd.to_datetime(raw_df['new_date'], format='%Y-%m-%d')

        # 应用截止日期过滤
        if cutoff_date:
            self.cutoff_date = pd.to_datetime(cutoff_date, format='%Y-%m-%d')
            raw_df = raw_df[raw_df['date_id'] <= self.cutoff_date]

        # 生成完整日期范围
        full_data = []
        grouped = raw_df.groupby('spu_code')

        for spu, group in grouped:
            launch_date = group['new_date'].min()

            # 计算商品生命周期内的日期范围
            date_range = pd.date_range(
                start=launch_date,
                periods=self.max_days,
                freq='D'
            )

            full_dates = pd.DataFrame({'date_id': date_range, 'spu_code': spu})

            # 合并实际销售数据
            merged = pd.merge(full_dates, group, on=['date_id', 'spu_code'], how='left')
            merged['daily_sale_qty'] = merged['daily_sale_qty'].fillna(0)
            # merged['days_since_launch'] = (merged['date_id'] - launch_date).dt.days
            merged['new_date'] = launch_date

            full_data.append(merged)

        full_df = pd.concat(full_data, ignore_index=True)

        # 生成训练样本
        X, y = [], []
        for spu, group in full_df.groupby('spu_code'):
            sales = group.sort_values('date_id')['daily_sale_qty'].values

            # 确保有90天的数据（不足补0）
            if len(sales) < self.max_days:
                sales = np.pad(sales, (0, self.max_days - len(sales)), 'constant')

            # 生成每个可能的预测点
            for start_day in range(1, self.max_days - self.forecast_days + 1):
                # 特征：历史销量 + 时间特征
                features = list(sales[:start_day])

                # 填充0使特征长度一致
                features += [0] * (self.max_days - len(features))

                # 添加时间特征
                features.append(start_day)  # 当前天数
                features.append(sum(sales[:start_day]))  # 累计销量
                features.append(np.mean(sales[:start_day]))  # 平均销量

                # 目标：未来14天累计销量
                end_day = min(start_day + self.forecast_days, len(sales))
                future_sales = sum(sales[start_day:end_day])

                X.append(features)
                y.append(future_sales)

        return np.array(X), np.array(y)

    def full_train(self, data_path, cutoff_date=None):
        """全量训练模型"""
        # 读取Excel文件
        if data_path.endswith('.xlsx'):
            raw_df = pd.read_excel(data_path, engine='openpyxl')
        elif data_path.endswith('.xls'):
            raw_df = pd.read_excel(data_path)
        else:
            raise ValueError("Unsupported file format. Only .xlsx and .xls files are supported.")



        X, y = self.prepare_data(raw_df, cutoff_date)

        # 检查是否有足够的数据
        if len(X) == 0:
            print("Warning: No training data available after filtering.")
            return None

        # 训练模型
        self.model = RandomForestRegressor(n_estimators=100, random_state=42)
        self.model.fit(X, y)

        # 保存模型
        joblib.dump(self.model, self.model_path)
        print(f"Model trained on {len(X)} samples and saved to {self.model_path}")
        return self.model

    def incremental_train(self, new_data_path):
        """增量训练模型"""
        # 加载现有模型
        try:
            self.model = joblib.load(self.model_path)
        except:
            print("No existing model found. Creating new model.")
            self.model = RandomForestRegressor(n_estimators=100, random_state=42)

        # 读取Excel文件
        if new_data_path.endswith('.xlsx'):
            new_df = pd.read_excel(new_data_path, engine='openpyxl')
        elif new_data_path.endswith('.xls'):
            new_df = pd.read_excel(new_data_path)
        else:
            raise ValueError("Unsupported file format. Only .xlsx and .xls files are supported.")

        # 准备增量数据
        X_new, y_new = self.prepare_data(new_df, self.cutoff_date)

        # 检查是否有新数据
        if len(X_new) == 0:
            print("Warning: No new training data available after filtering.")
            return None

        # 增量训练
        self.model.n_estimators += 10  # 增加树的数量
        self.model.fit(X_new, y_new)

        # 保存更新后的模型
        joblib.dump(self.model, self.model_path)
        print(f"Model incrementally trained with {len(X_new)} new samples")
        return self.model

    def predict(self, spu_data):
        """预测单个SPU未来14天销量"""
        # 确保模型已加载
        if self.model is None:
            try:
                self.model = joblib.load(self.model_path)
            except:
                raise ValueError("Model not trained or loaded. Please train a model first.")

        # 确保输入是DataFrame
        if not isinstance(spu_data, pd.DataFrame):
            raise ValueError("Input must be a pandas DataFrame")

        # 转换日期格式
        spu_data['date_id'] = pd.to_datetime(spu_data['date_id'], format='%Y-%m-%d')
        spu_data['new_date'] = pd.to_datetime(spu_data['new_date'], format='%Y-%m-%d')

        # 获取最新状态
        last_date = spu_data['date_id'].max()
        launch_date = spu_data['new_date'].min()

        # 计算当前天数
        current_day = (last_date - launch_date).days

        # 检查是否超过预测范围
        if current_day >= self.max_days - self.forecast_days:
            spu_code = spu_data['spu_code'].iloc[0] if 'spu_code' in spu_data.columns else 'Unknown'
            print(f"SPU {spu_code} exceeds 90-day lifecycle (current day: {current_day}), no prediction made")
            return 0
        if current_day < 0:
            spu_code = spu_data['spu_code'].iloc[0] if 'spu_code' in spu_data.columns else 'Unknown'
            print(f"Invalid date range for SPU {spu_code} (launch date: {launch_date}, last date: {last_date})")
            return 0

        # 准备特征向量
        sales = spu_data['daily_sale_qty'].values

        # 确保销售数据长度足够
        features = list(sales)
        if len(features) < self.max_days:
            features += [0] * (self.max_days - len(features))
        else:
            features = features[:self.max_days]

        # 添加时间特征
        features.append(current_day)
        # print(sales)
        features.append(sum(sales))
        features.append(np.mean(sales) if len(sales) > 0 else 0)

        # 预测
        try:
            prediction = self.model.predict([features])[0]
            return max(0, prediction)  # 确保非负预测
        except Exception as e:
            print(f"Prediction failed: {str(e)}")
            return 0


# 示例用法
if __name__ == "__main__":
    predictor = SalesPredictor()

    # 全量训练（使用截止日期）
    predictor.full_train('sales.xlsx', cutoff_date='2025-06-01')

    # 增量训练（添加新一天的数据）
    # predictor.incremental_train('new_day_sales.xlsx')

    # 准备预测数据示例



    # new_spu_data = pd.DataFrame({
    #     'spu_code': ['LO2505011','LO2505011','LO2505011','LO2505011','LO2505011','LO2505011','LO2505011','LO2505011','LO2505011','LO2505011'],
    #     'date_id': ['2025-05-31','2025-06-01','2025-06-02','2025-06-03','2025-06-04','2025-06-05','2025-06-06','2025-06-07','2025-06-08','2025-06-09'],
    #     'new_date': ['2025-05-31','2025-05-31','2025-05-31','2025-05-31','2025-05-31','2025-05-31','2025-05-31','2025-05-31','2025-05-31','2025-05-31'],
    #     'daily_sale_qty': [1028,151,91,109,135,89,56,83,90,58],
    #     'days_since_launch': [0,1,2,3,4,5,6,7,8,9]
    # })

    known_days = 15

    new_spu_data = pd.DataFrame({
        'spu_code': ['LO2505011']*known_days,
        'date_id': pd.date_range(start='2025-05-31', periods=known_days, freq='D'),
        'new_date': ['2025-05-31']*known_days,
        'daily_sale_qty': [1028,151,91,109,135,89,56,83,90,58,72,59,55,40,59],
        # 'days_since_launch': [0,1,2,3,4,5,6,7,8,9,10,11,12,13,14]
    })


    # new_spu_data = pd.DataFrame({
    #     'spu_code': ['LO2505020'],
    #     'date_id': ['2025-05-31'],
    #     'new_date': ['2025-05-31'],
    #     'daily_sale_qty': [370],
    #     'days_since_launch': [0]
    # })

    # new_spu_data = pd.DataFrame({
    #     'spu_code': ['LO2505020','LO2505020'],
    #     'date_id': ['2025-05-31','2025-06-01'],
    #     'new_date': ['2025-05-31','2025-05-31'],
    #     'daily_sale_qty': [370,37],
    #     'days_since_launch': [0,1]
    # })

    # 执行预测
    forecast = predictor.predict(new_spu_data)
    print(f"Predicted 14-day sales: {forecast:.2f}")